/*************************************************************
 * TwoSocks for network connection toy.
 *
 * This software utilize two sockets in modes of server and
 * client. Actually two sockets handle two connections.
 * 
 * (C) Copyright by Hojin Choi <hojin.choi@gmail.com>
 *
 * License: Apache License v2.0
 * Last Date: 2010.07.11
 * Version: 0.1
 *************************************************************/

#include "config.h"
#include "2socks.h"
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

#if HAVE_UNISTD_H
#include <unistd.h>
#else
#error "We need unistd.h for using getopt"
#endif

#include <string.h>

#if HAVE_ARPA_INET_H
#include <arpa/inet.h>
#endif

#if HAVE_NETDB_H
#include <netdb.h>
#endif

struct global_options ts_opts;

static struct sockaddr ts_get_bindaddr( int family, const char * addr, const char * portstr )
{
	struct sockaddr bindaddr;
	struct sockaddr_in  * p4 = (struct sockaddr_in  *) &bindaddr;
	struct sockaddr_in6 * p6 = (struct sockaddr_in6 *) &bindaddr;

	switch( family ) {
	case AF_INET:
		if( 1 != inet_pton( family, addr, & p4->sin_addr ) ) {
			memset( &bindaddr, 0, sizeof bindaddr );
		}
		p4->sin_port = htons( atoi( portstr ) );
		break;
	case AF_INET6:
		if( 1 != inet_pton( family, addr, & p6->sin6_addr ) ) {
			memset( &bindaddr, 0, sizeof bindaddr );
		}
		p6->sin6_port = htons( atoi( portstr ) );
		break;
	default:
		printf("Not implemented yet.. address family(%d)\n", family);
		abort();
	}
	bindaddr.sa_family = family;
	return bindaddr;
}

static struct addrinfo * ts_gethostaddr( int family, const char * host, const char * portstr )
{
	struct addrinfo hints;
	struct addrinfo *result, *rp;
	int port;
	int err;

	memset(&hints, 0, sizeof(struct addrinfo));
	hints.ai_family = family;
	hints.ai_socktype = SOCK_DGRAM;
	hints.ai_flags = AI_PASSIVE;
	hints.ai_protocol = 0;
	hints.ai_canonname = NULL;
	hints.ai_addr = NULL;
	hints.ai_next = NULL;

	err = getaddrinfo(host, NULL, &hints, &result);

	if (err != 0) {
		fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(err));
		return 0;
	}

	port = htons( atoi( portstr ) );
	for (rp = result; rp != NULL; rp = rp->ai_next) {
		if( rp->ai_family == AF_INET ) {
			((struct sockaddr_in *) rp->ai_addr)->sin_port = port;
		}
		if( rp->ai_family == AF_INET6 ) {
			((struct sockaddr_in6 *) rp->ai_addr)->sin6_port = port;
		}
	}

	return result;
}

static void init_opts( struct options * popts )
{
	memset( popts, 0, sizeof *popts );
}

static int parse_fd( struct ts_address * node, const char * arg )
{
	node->family = AF_PRIVATE_STDIO;
	if( 0 == strcmp( arg, "stdio" ) )
	{
		node->inherited_fd = fileno(stdin);
		return 0;
	}
	node->inherited_fd = atoi(arg);
	if( arg[0] != '0' && node->inherited_fd == 0 ) {
		ts_exit( 1, "Invalid file descriptor" );
	}
	return 0;
}

static int parse_ipv4( struct ts_address * node, char * const parts[], int count, int isserver )
{
	node->family = AF_INET;
	node->len_bindaddr = sizeof( struct sockaddr_in );

	switch( count )
	{
	case 1:
		if( isserver ) {
			node->bindaddr = ts_get_bindaddr( AF_INET, "", parts[0] );
		} else {
			node->addr = ts_gethostaddr( PF_INET, "localhost", parts[0] );
		}
		break;
	case 2:
		if( isserver ) {
			node->bindaddr = ts_get_bindaddr( AF_INET, parts[0], parts[1] );
		} else {
			node->addr = ts_gethostaddr( PF_INET, parts[0], parts[1] );
		}
		break;
	case 3:
		if( isserver ) {
			/* Ignore parts 0 */
			node->bindaddr = ts_get_bindaddr( AF_INET, parts[1], parts[2] );
		} else {
			node->bindaddr = ts_get_bindaddr( AF_INET, "", parts[0] );
			node->addr = ts_gethostaddr( PF_INET, parts[1], parts[2] );
		}
		break;
	case 4:
		if( isserver ) {
			/* Ignore parts 0, 1 */
			node->bindaddr = ts_get_bindaddr( AF_INET, parts[2], parts[3] );
		} else {
			node->bindaddr = ts_get_bindaddr( AF_INET, parts[0], parts[1] );
			node->addr = ts_gethostaddr( PF_INET, parts[2], parts[3] );
		}
		break;
	}
	if( !isserver && !node->addr )
	{
		exit(0);
	}
	return 0;
}

static int parse_arg( int isserver, const char * arg, struct ts_address * node )
{
	char * strctx;
	char * part;
	char * temp = strdup( arg );
	char * parts[4] = {"","","",""}; /* 4 means maximum parts */
	int count = 0;

	assert( arg );
	assert( node );

	if( NULL == temp )
	{
		return -1;
	}

	for( part = strtok_r( temp, ":", &strctx );
	     part && count < 4;
	     part = strtok_r( NULL, ":", &strctx ) )
	{
		/* printf("TOK %d:%s\n", count, part ); */
		parts[count] = strdup(part);
		count++;
	}
	free( temp );

	memset( node, 0, sizeof * node );

	if( count > 1 )
	{
		if( strcmp( parts[0], "fd" ) == 0 )
		{
			if( isserver ) {
				return -1;
			}
			return parse_fd( node, parts[1] );
		}
	}

	/* TCP/IP */
	return parse_ipv4( node, parts, count, isserver );
}

void usage()
{
	puts( "Usage: 2socks [-vh] [-sc|-ss|-cc|cs] <args> [<args>...]" );
	puts( "   Two sockets are created and they are used as lego blocks" );
	if( ! ts_opts.verbosity ) {
		puts( "   -sc <listen port> [<remote address>:]<target port>" );
		puts( "   -ss <listen port> <listen port>" );
		puts( "   -cc [<remote address>:]<target port> [<remote address>:]<target port>" );
		/* puts( "   -cs [<remote address>:]<target port> <listen port>" ); */
	} else {
		puts( "   -sc [<bind address>:]<listen port> [[[<bind address>:]<bind port>:]<remote address>:]<target port>" );
		puts( "   -ss [<bind address>:]<listen port> [<bind address>:]<listen port>" );
		puts( "   -cc [[<bind port>:]<remote address>:]<target port> [[[<bind address>:]<bind port>:]<remote address>:]<target port>" );
		/* puts( "   -cs [[<bind port>:]<remote address>:]<target port> [<bind address>:]<listen port>" ); */
	}
	puts( "" );
	puts( "Mode: -sc; Serving and then Connect mode\n"
	        "    This mode creates a server socket, and after accepting connection\n"
	        "    a socket is created and tried to connect remote address specified second argument\n"
	        "    for relaying the first connection\n"
	        "    Default remote address is localhost\n"
	        "    You can specify client binding port for connecting remote host" );
	puts( "" );
	puts( "Mode: -ss; Two Servers Mode" );
	puts( "" );
	puts( "Mode: -cc; Two Clients mode" );
	puts( "" );
	/*
	printf( "Mode: -cs; Connect and then Serving mode" );
	printf( "" );
	*/
	puts( "Options:" );
	puts( "   -v : Print messages verbosely" );
	puts( "   -h : Print this help message (-vh: print long arguments)" );


	exit(0);
}

#define CHECK_UNIQ_ENGINE() do { if( engine_type != TS_NONE ) { \
	fprintf( stderr, "You should specify only one type at once.\n" ); \
	exit(1);} } while(0)

static int get_engine_type( struct options * opts, int argc, char ** argv )
{
	int ch;
	int engine_type = TS_NONE;
	void (*engine)( const struct options * );

	while ((ch = getopt(argc, argv, "hvs:c:")) != -1) {
		switch (ch) {
		case 's':
			if( *optarg == 's' )
			{
				CHECK_UNIQ_ENGINE();
				engine_type = TS_SERVER_SERVER;
			}
			else if( *optarg == 'c' )
			{
				CHECK_UNIQ_ENGINE();
				engine_type = TS_SERVER_CLIENT;
			}
			break;
		case 'c':
			if( *optarg == 's' )
			{
				CHECK_UNIQ_ENGINE();
				engine_type = TS_CLIENT_SERVER;
			}
			else if( *optarg == 'c' )
			{
				CHECK_UNIQ_ENGINE();
				engine_type = TS_CLIENT_CLIENT;
			}
			break;
		case 'v':
			ts_opts.verbosity = 1;
			break;
		case 'h':
		default:
			usage();
		}
	}

	if( argv[optind+0] == NULL || argv[optind+1] == NULL )
	{
		usage();
	}

	opts->args[0] = argv[optind+0];
	opts->args[1] = argv[optind+1];

	return engine_type;
}

#define ipport(a) \
	ntohs( \
		(((struct sockaddr *)&a)->sa_family == AF_INET ) ? \
			((struct sockaddr_in*) &a)->sin_port : \
			((struct sockaddr_in6*) &a)->sin6_port )

void check_addr_sanity( int family, struct sockaddr * addr )
{
	if( family == AF_PRIVATE_STDIO ) {
		return;
	}
	if( ipport(addr) == 0 ) {
		usage();
	}
}

#define CHECK_SANITY( node, addr ) check_addr_sanity( node.family, (struct sockaddr *) & node.addr )

static engine_t get_engine( int engine_type, struct options * opts )
{
	engine_t engine;
	switch( engine_type )
	{
	case TS_SERVER_SERVER:
		engine = server_server;
		parse_arg( 1, opts->args[0], &opts->nodes[TS_FIRST] );
		parse_arg( 1, opts->args[1], &opts->nodes[TS_SECOND] );
		CHECK_SANITY( opts->nodes[TS_FIRST], bindaddr );
		CHECK_SANITY( opts->nodes[TS_SECOND], bindaddr );
		break;

	case TS_SERVER_CLIENT:
		engine = server_client;
		parse_arg( 1, opts->args[0], &opts->nodes[TS_FIRST] );
		parse_arg( 0, opts->args[1], &opts->nodes[TS_SECOND] );
		CHECK_SANITY( opts->nodes[TS_FIRST], bindaddr );
		CHECK_SANITY( opts->nodes[TS_SECOND], addr->ai_addr );
		break;

	case TS_CLIENT_SERVER:
		engine = client_server;
		parse_arg( 0, opts->args[0], &opts->nodes[TS_FIRST] );
		parse_arg( 1, opts->args[1], &opts->nodes[TS_SECOND] );
		break;

	case TS_CLIENT_CLIENT:
		engine = client_client;
		parse_arg( 0, opts->args[0], &opts->nodes[TS_FIRST] );
		parse_arg( 0, opts->args[1], &opts->nodes[TS_SECOND] );
		CHECK_SANITY( opts->nodes[TS_FIRST], addr->ai_addr );
		CHECK_SANITY( opts->nodes[TS_SECOND], addr->ai_addr );
		break;
	default:
		usage();
		break;
	}

	return engine;
}

int main( int argc, char ** argv )
{
	struct options opts;
	engine_t engine;
	int engine_type = TS_NONE;

	engine = NULL;
	init_opts( & opts );

	engine_type = get_engine_type( & opts, argc, argv );
	engine = get_engine( engine_type, & opts );

	if( engine )
	{
		engine( & opts );
	}
	return 0;
}
